import numpy as np
from scipy.spatial import cKDTree as KDTree
import cli

# Checking
from shapely.geometry import Point, Polygon

def sph_to_cart(lon, lat):
	R = 6371e3 # m
	x = R * np.cos(lat) * np.cos(lon)
	y = R * np.cos(lat) * np.sin(lon)
	z = R * np.sin(lat)
	return (x, y, z)

def create_kdtree(xlon, xlat):
	rxlon = np.radians(xlon)
	rxlat = np.radians(xlat)
	imap = np.ndarray(xlon.size, dtype='i4')
	jmap = np.ndarray(xlon.size, dtype='i4')
	k = 0
	for j in range(xlon.shape[1]):
		for i in range(xlon.shape[0]):
			imap[k] = i
			jmap[k] = j
			k += 1
	X = np.transpose(np.array(sph_to_cart(rxlon, rxlat))).reshape((xlon.size, 3))
	return (KDTree(X), imap, jmap, xlon, xlat)

def find_bbox(kdtree, lon, lat):
	res = kdtree[0].query(sph_to_cart(np.radians(lon), np.radians(lat)))
	j0 = kdtree[1][res[1]]
	i0 = kdtree[2][res[1]]
	lon0 = float(kdtree[3][j0,i0])
	lat0 = float(kdtree[4][j0,i0])
	bbox = np.ndarray((4,2), dtype='i4')
	if lon >= lon0:
		bbox[0,1] = i0
		bbox[1,1] = i0 + 1
		bbox[2,1] = i0 + 1
		bbox[3,1] = i0
	else:
		bbox[0,1] = i0 - 1
		bbox[1,1] = i0
		bbox[2,1] = i0
		bbox[3,1] = i0 - 1
	if lat >= lat0:
		bbox[0,0] = j0
		bbox[1,0] = j0
		bbox[2,0] = j0 + 1
		bbox[3,0] = j0 + 1
	else:
		bbox[0,0] = j0 - 1
		bbox[1,0] = j0 - 1
		bbox[2,0] = j0
		bbox[3,0] = j0
	if np.any(bbox[:,0] < 0) or np.any(bbox[:,0] >= kdtree[3].shape[0]) or np.any(bbox[:,1] < 0) or np.any(bbox[:,1] >= kdtree[4].shape[1]):
		cli.error('Point is out of domain!')

	# pt = Point(lon, lat)
	# poly = Polygon([
	# 	[float(kdtree[3][bbox[0,0],bbox[0,1]]),float(kdtree[4][bbox[0,0],bbox[0,1]])],
	# 	[float(kdtree[3][bbox[1,0],bbox[1,1]]),float(kdtree[4][bbox[1,0],bbox[1,1]])],
	# 	[float(kdtree[3][bbox[2,0],bbox[2,1]]),float(kdtree[4][bbox[2,0],bbox[2,1]])],
	# 	[float(kdtree[3][bbox[3,0],bbox[3,1]]),float(kdtree[4][bbox[3,0],bbox[3,1]])],
	# 	[float(kdtree[3][bbox[0,0],bbox[0,1]]),float(kdtree[4][bbox[0,0],bbox[0,1]])]
	# ])
	# if not poly.contains(pt):
	#	cli.error('Point is not in polygon!')
	return bbox

def create_bilinear_wgts(xlon, xlat, bbox, lon, lat):
	jab = np.ndarray((2, 2))
	lon0 = xlon[bbox[0,0],bbox[0,1]]; lat0 = xlat[bbox[0,0],bbox[0,1]]
	lon1 = xlon[bbox[1,0],bbox[1,1]]; lat1 = xlat[bbox[1,0],bbox[1,1]]
	lon2 = xlon[bbox[2,0],bbox[2,1]]; lat2 = xlat[bbox[2,0],bbox[2,1]]
	lon3 = xlon[bbox[3,0],bbox[3,1]]; lat3 = xlat[bbox[3,0],bbox[3,1]]
	jab[0,0] = ((lon1 + lon2) - (lon3 + lon0)) * 0.5
	jab[0,1] = ((lon2 + lon3) - (lon0 + lon1)) * 0.5
	jab[1,0] = ((lat1 + lat2) - (lat3 + lat0)) * 0.5
	jab[1,1] = ((lat2 + lat3) - (lat0 + lat1)) * 0.5
	a, b = np.matmul(jab, np.array([lon, lat])) - np.matmul(jab, np.array([lon0,lat0]))
	return [np.array([(1 - a) * (1 - b), a * (1 - b), a * b, (1 - a) * b]), bbox]

def apply_bilinear(wgts, var, dims, debug=False):
	bbox = wgts[1]
	return np.sum(np.array([wgts[0][i] * var.isel({dims[0]: bbox[i,0],dims[1]: bbox[i,1]}) for i in range(4)]), axis=0)
